"""SSH Connection Management Module"""

import asyncio
import logging
import os
import tempfile
from pathlib import Path
from typing import Dict, Optional, Any, List, Tuple
import paramiko
from paramiko import SSHClient, SFTPClient
import time

logger = logging.getLogger(__name__)


class SSHManager:
    """Manages SSH connections and operations for remote servers"""
    
    def __init__(self):
        self.clients: Dict[str, SSHClient] = {}
        self.sftp_clients: Dict[str, SFTPClient] = {}
        self.connection_info: Dict[str, Dict[str, Any]] = {}
        
    def connect(self, host: str, username: str, key_path: Optional[str] = None, 
                password: Optional[str] = None, port: int = 22) -> str:
        """Establish SSH connection to remote server
        
        Args:
            host: Remote server hostname or IP
            username: SSH username
            key_path: Path to SSH private key file
            password: SSH password (alternative to key_path)
            port: SSH port (default: 22)
            
        Returns:
            Connection ID for future operations
            
        Raises:
            Exception: If connection fails
        """
        connection_id = f"{username}@{host}:{port}"
        
        try:
            # Create SSH client
            client = SSHClient()
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Prepare connection arguments
            connect_kwargs = {
                'hostname': host,
                'username': username,
                'port': port,
                'timeout': 10
            }
            
            # Use key-based authentication if key_path provided
            if key_path and os.path.exists(key_path):
                connect_kwargs['key_filename'] = key_path
            elif password:
                connect_kwargs['password'] = password
            else:
                raise ValueError("Either key_path or password must be provided")
            
            # Establish connection
            client.connect(**connect_kwargs)
            
            # Test connection
            _, stdout, stderr = client.exec_command('echo "Connection test"')
            stdout.read()
            
            # Create SFTP client
            sftp = client.open_sftp()
            
            # Store clients and connection info
            self.clients[connection_id] = client
            self.sftp_clients[connection_id] = sftp
            self.connection_info[connection_id] = {
                'host': host,
                'username': username,
                'port': port,
                'connected_at': time.time()
            }
            
            logger.info(f"Successfully connected to {connection_id}")
            return connection_id
            
        except Exception as e:
            logger.error(f"Failed to connect to {host}: {str(e)}")
            raise Exception(f"SSH connection failed: {str(e)}")
    
    def disconnect(self, connection_id: Optional[str] = None) -> bool:
        """Disconnect from SSH server
        
        Args:
            connection_id: Specific connection to disconnect (None = disconnect all)
            
        Returns:
            True if successful
        """
        try:
            if connection_id:
                # Disconnect specific connection
                if connection_id in self.clients:
                    self.sftp_clients[connection_id].close()
                    self.clients[connection_id].close()
                    del self.sftp_clients[connection_id]
                    del self.clients[connection_id]
                    del self.connection_info[connection_id]
                    logger.info(f"Disconnected from {connection_id}")
                    return True
                else:
                    logger.warning(f"Connection {connection_id} not found")
                    return False
            else:
                # Disconnect all connections
                for conn_id in list(self.clients.keys()):
                    self.disconnect(conn_id)
                logger.info("Disconnected from all SSH connections")
                return True
                
        except Exception as e:
            logger.error(f"Error disconnecting: {str(e)}")
            return False
    
    def get_client(self, connection_id: str) -> SSHClient:
        """Get SSH client for connection
        
        Args:
            connection_id: Connection identifier
            
        Returns:
            SSH client instance
            
        Raises:
            ValueError: If connection not found
        """
        if connection_id not in self.clients:
            raise ValueError(f"No active connection found: {connection_id}")
        return self.clients[connection_id]
    
    def get_sftp_client(self, connection_id: str) -> SFTPClient:
        """Get SFTP client for connection
        
        Args:
            connection_id: Connection identifier
            
        Returns:
            SFTP client instance
            
        Raises:
            ValueError: If connection not found
        """
        if connection_id not in self.sftp_clients:
            raise ValueError(f"No active SFTP connection found: {connection_id}")
        return self.sftp_clients[connection_id]
    
    def execute_command(self, connection_id: str, command: str, 
                       working_dir: str = "~") -> Tuple[int, str, str]:
        """Execute command on remote server
        
        Args:
            connection_id: SSH connection identifier
            command: Command to execute
            working_dir: Working directory for command execution
            
        Returns:
            Tuple of (exit_code, stdout, stderr)
            
        Raises:
            ValueError: If connection not found
        """
        client = self.get_client(connection_id)
        
        try:
            # Change to working directory and execute command
            full_command = f"cd {working_dir} && {command}"
            stdin, stdout, stderr = client.exec_command(full_command)
            
            # Wait for command completion and get results
            exit_code = stdout.channel.recv_exit_status()
            stdout_data = stdout.read().decode('utf-8')
            stderr_data = stderr.read().decode('utf-8')
            
            logger.info(f"Command executed on {connection_id}: {command}")
            return exit_code, stdout_data, stderr_data
            
        except Exception as e:
            logger.error(f"Error executing command: {str(e)}")
            raise Exception(f"Command execution failed: {str(e)}")
    
    def upload_file(self, connection_id: str, local_path: str, remote_path: str) -> bool:
        """Upload file from local to remote server
        
        Args:
            connection_id: SSH connection identifier
            local_path: Local file path
            remote_path: Remote destination path
            
        Returns:
            True if successful
            
        Raises:
            ValueError: If connection not found or file issues
        """
        if not os.path.exists(local_path):
            raise ValueError(f"Local file not found: {local_path}")
        
        sftp = self.get_sftp_client(connection_id)
        
        try:
            # Ensure remote directory exists
            remote_dir = os.path.dirname(remote_path)
            if remote_dir and remote_dir != '/':
                try:
                    sftp.mkdir(remote_dir)
                except OSError:
                    pass  # Directory might already exist
            
            # Upload file
            sftp.put(local_path, remote_path)
            logger.info(f"Uploaded {local_path} to {connection_id}:{remote_path}")
            return True
            
        except Exception as e:
            logger.error(f"Error uploading file: {str(e)}")
            raise Exception(f"File upload failed: {str(e)}")
    
    def download_file(self, connection_id: str, remote_path: str, local_path: str) -> bool:
        """Download file from remote server to local
        
        Args:
            connection_id: SSH connection identifier
            remote_path: Remote file path
            local_path: Local destination path
            
        Returns:
            True if successful
            
        Raises:
            ValueError: If connection not found
        """
        sftp = self.get_sftp_client(connection_id)
        
        try:
            # Ensure local directory exists
            local_dir = os.path.dirname(local_path)
            if local_dir:
                os.makedirs(local_dir, exist_ok=True)
            
            # Download file
            sftp.get(remote_path, local_path)
            logger.info(f"Downloaded {connection_id}:{remote_path} to {local_path}")
            return True
            
        except Exception as e:
            logger.error(f"Error downloading file: {str(e)}")
            raise Exception(f"File download failed: {str(e)}")
    
    def list_remote_files(self, connection_id: str, directory: str = "~", 
                         pattern: str = "*") -> List[Dict[str, Any]]:
        """List files in remote directory
        
        Args:
            connection_id: SSH connection identifier
            directory: Remote directory path
            pattern: File pattern filter
            
        Returns:
            List of file information dictionaries
            
        Raises:
            ValueError: If connection not found
        """
        try:
            # Use ls command for better pattern support
            exit_code, stdout, stderr = self.execute_command(
                connection_id, f"ls -la {directory}/{pattern}", working_dir="~"
            )
            
            if exit_code != 0:
                raise Exception(f"ls command failed: {stderr}")
            
            files = []
            for line in stdout.strip().split('\n'):
                if line and not line.startswith('total'):
                    parts = line.split(None, 8)
                    if len(parts) >= 9:
                        files.append({
                            'permissions': parts[0],
                            'links': parts[1],
                            'owner': parts[2],
                            'group': parts[3],
                            'size': parts[4],
                            'date': f"{parts[5]} {parts[6]} {parts[7]}",
                            'name': parts[8],
                            'full_line': line
                        })
            
            logger.info(f"Listed {len(files)} files in {connection_id}:{directory}")
            return files
            
        except Exception as e:
            logger.error(f"Error listing remote files: {str(e)}")
            raise Exception(f"Remote file listing failed: {str(e)}")
    
    def is_connected(self, connection_id: str) -> bool:
        """Check if connection is active
        
        Args:
            connection_id: Connection identifier
            
        Returns:
            True if connection is active
        """
        if connection_id not in self.clients:
            return False
        
        try:
            client = self.clients[connection_id]
            transport = client.get_transport()
            return transport is not None and transport.is_active()
        except:
            return False
    
    def list_connections(self) -> List[Dict[str, Any]]:
        """List all active connections
        
        Returns:
            List of connection information
        """
        connections = []
        for conn_id, info in self.connection_info.items():
            connections.append({
                'connection_id': conn_id,
                'host': info['host'],
                'username': info['username'],
                'port': info['port'],
                'connected_at': info['connected_at'],
                'is_active': self.is_connected(conn_id)
            })
        return connections
    
    def cleanup(self):
        """Clean up all connections"""
        self.disconnect()
